﻿#ifndef AVLTREE_H
#define AVLTREE_H
#include <cassert>
#include <iostream>
#include <print>
using namespace std;

template<class K, class V>
struct AVLTreeNode
{
    pair<K, V> kv_;       // 存储键值对
    AVLTreeNode *left_;   // 左子树节点指针
    AVLTreeNode *right_;  // 右子树节点指针
    AVLTreeNode *parent_; // 父节点指针
    int balanceFactor_;   // 平衡因子

    explicit AVLTreeNode(const pair<K, V> &kv)
        : kv_(kv)
        , left_(nullptr)
        , right_(nullptr)
        , parent_(nullptr)
        , balanceFactor_(0)
    {}
};

template<class K, class V>
class AVLTree
{
    using Node = AVLTreeNode<K, V>;

public:
    void in_order()
    {
        in_order_(root_);
        println();
    }

    bool insert(const pair<K, V> &kv)
    {
        if (root_ == nullptr)
        {
            root_ = new Node(kv);
            return true;
        }

        Node *parent = nullptr;
        Node *cur = root_;
        while (cur)
        {
            if (cur->kv_.first < kv.first)
            {
                parent = cur;
                cur = cur->right_;
            }
            else if (cur->kv_.first > kv.first)
            {
                parent = cur;
                cur = cur->left_;
            }
            else
            {
                return false;
            }
        }

        cur = new Node(kv);
        cur->parent_ = parent;
        if (parent->kv_.first < kv.first)
        {
            parent->right_ = cur;
        }
        else
        {
            parent->left_ = cur;
        }

        while (parent)
        {
            if (cur == parent->right_)
            {
                ++parent->balanceFactor_;
            }
            else
            {
                --parent->balanceFactor_;
            }
            if (parent->balanceFactor_ == 0)
            {
                break;
            }
            else if (parent->balanceFactor_ == 1 || parent->balanceFactor_ == -1)
            {
                cur = parent;
                parent = parent->parent_;
            }
            else if (parent->balanceFactor_ == 2 || parent->balanceFactor_ == -2)
            {
                if (parent->balanceFactor_ == -2 && cur->balanceFactor_ == -1)
                    rotate_right(parent);
                else if (parent->balanceFactor_ == 2 && cur->balanceFactor_ == 1)
                    rotate_left(parent);
                else if (parent->balanceFactor_ == -2 && cur->balanceFactor_ == 1)
                    rotate_left_right(parent);
                else if (parent->balanceFactor_ == 2 && cur->balanceFactor_ - 1)
                    rotate_right_left(parent);
                else
                    assert(false);
                break;
            }
            else
            {
                assert(false);
            }
        }

        return true;
    }
    bool is_balance()
    {
        return is_balance_(root_);
    }

private:
    int height_(Node *root)
    {
        if (root == nullptr)
            return 0;
        const int leftHeight = height_(root->left_);
        const int rightHeight = height_(root->right_);
        return leftHeight > rightHeight ? leftHeight + 1 : rightHeight + 1;
    }
    bool is_balance_(Node *root)
    {
        if (root == nullptr)
            return true;
        const int leftHeight = height_(root->left_);
        const int rightHeight = height_(root->right_);
        const int balanceFactor = rightHeight - leftHeight;

        if (abs(balanceFactor) >= 2 || balanceFactor != root->balanceFactor_)
        {
            println("{}, 平衡因子有问题", root->kv_.first);
            return false;
        }

        return is_balance_(root->left_) && is_balance_((root->right_));
    }
    void rotate_right(Node *parent)
    {
        Node *subtreeLeft = parent->left_;
        Node *subtreeRight = parent->right_;

        parent->left_ = subtreeRight;
        if (subtreeLeft)
            subtreeRight->parent_ = parent;

        Node *parentParent = parent->parent_;

        subtreeLeft->right_ = parent;
        parent->parent_ = subtreeLeft;

        if (parent == root_)
        {
            root_ = subtreeLeft;
            subtreeLeft->parent_ = nullptr;
        }
        else
        {
            if (parentParent->left_ == parent)
            {
                parentParent->left_ = subtreeLeft;
            }
            else
            {
                parentParent->right_ = subtreeLeft;
            }
            parent->balanceFactor_ = subtreeLeft->balanceFactor_;
        }
        parent->balanceFactor_ = subtreeLeft->balanceFactor_ = 0;
    }

    void rotate_left(Node *parent)
    {
        Node *subtreeRight = parent->right_;
        Node *subtreeRightLeft = parent->left_;

        parent->right_ = subtreeRightLeft;
        if (subtreeRightLeft)
            subtreeRightLeft->parent_ = parent;

        Node *parentParent = parent->parent_;

        subtreeRight->left_ = parent;
        parent->parent_ = subtreeRight;

        if (parent == root_)
        {
            root_ = subtreeRight;
            subtreeRight->parent_ = nullptr;
        }
        else
        {
            if (parentParent->left_ == parent)
            {
                parentParent->left_ = subtreeRight;
            }
            else
            {
                parentParent->right_ = subtreeRight;
            }
            parent->balanceFactor_ = subtreeRight->balanceFactor_;
        }
        parent->balanceFactor_ = subtreeRight->balanceFactor_ = 0;
    }
    void rotate_left_right(Node *parent)
    {
        Node *subtreeLeft = parent->left_;
        Node *subtreeLeftRight = subtreeLeft->right_;

        const int balanceFactor = subtreeLeft->balanceFactor_;

        rotate_left(subtreeLeft);
        rotate_right(subtreeLeftRight);

        if (balanceFactor == 0)
        {
            subtreeLeft->balanceFactor_ = 0;
            parent->balanceFactor_ = 0;
            subtreeLeftRight->balanceFactor_ = 0;
        }
        else if (balanceFactor == 1)
        {
            subtreeLeft->balanceFactor_ = -1;
            parent->balanceFactor_ = 0;
            subtreeLeftRight->balanceFactor_ = 0;
        }
        else if (balanceFactor == -1)
        {
            subtreeLeft->balanceFactor_ = 0;
            parent->balanceFactor_ = 1;
            subtreeLeftRight->balanceFactor_ = 0;
        }
        else
        {
            assert(false);
        }
    }
    void rotate_right_left(Node *parent)
    {
        Node *subtreeRight = parent->right_;
        Node *subtreeRightLeft = subtreeRight->left_;

        const int balanceFactor = subtreeRightLeft->balanceFactor_;

        rotate_right(parent->right_);
        rotate_left(parent->left_);

        if (balanceFactor == 0)
        {
            subtreeRight->balanceFactor_ = 0;
            subtreeRightLeft->balanceFactor_ = 0;
            parent->balanceFactor_ = 0;
        }
        else if (balanceFactor == 1)
        {
            subtreeRight->balanceFactor_ = 0;
            subtreeRightLeft->balanceFactor_ = 0;
            parent->balanceFactor_ = -1;
        }
        else if (balanceFactor == -1)
        {
            subtreeRight->balanceFactor_ = 1;
            subtreeRightLeft->balanceFactor_ = 0;
            parent->balanceFactor_ = 0;
        }
        else
        {
            assert(false);
        }
    }

    void in_order_(Node *root)
    {
        if (root == nullptr)
            return;
        in_order_(root->left_);
        print("{} ", root->kv_.first);
        // println("{}:{}", root->kv_.first, root->kv_.second);
        in_order_(root->right_);
    }

    Node *root_ = nullptr;
};
#endif /* AVLTREE_H */
